import os
import random
from random import shuffle

from PIL import Image
from torch.utils.data import Dataset

from config import parser

args = parser.parse_args()
from torchvision import transforms

train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomVerticalFlip(),
    # transforms.RandomCrop(50),
    transforms.ColorJitter(brightness=0.5, contrast=0.5, hue=0.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

])
test_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    # transforms.RandomVerticalFlip(),
    # transforms.RandomCrop(50),
    # transforms.ColorJitter(brightness=0.5, contrast=0.5, hue=0.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

])


# 自己写Dataset至少需要有这样的格式
class Dataset(Dataset):
    def __init__(self, type):
        super(Dataset, self).__init__()

        self.type = "train" if type == "train" else "val"
        self.base_path = args.base_data_path + self.type + '/'
        self.imgs = os.listdir(self.base_path + '/cat') + os.listdir(self.base_path + '/dog')
        random.shuffle(self.imgs)

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, index):
        if index == 0:
            shuffle(self.imgs)
        n = len(self.imgs)
        index = index % n
        img, y = self.collect_image_label(self.imgs[index])
        img = img.resize((448, 448), Image.BICUBIC)
        if self.type == "train":
            temp_img = train_transform(img)
        else:
            temp_img = test_transform(img)
        temp_y = int(y)
        return temp_img, temp_y

    def collect_image_label(self, line):
        name = line.split('.')[0]
        image_path = self.base_path + '/' + name + '/' + line

        if name == 'dog':
            label = 0
        else:
            label = 1

        image = Image.open(image_path).convert("RGB")
        # image=Image.new(mode='RGB',size=(448,448),color=(0,0,0))

        return image, label

    # # DataLoader中collate_fn使用
    # def dataset_collate(batch):
    #     images = []
    #     bboxes = []
    #     for img, box in batch:
    #         images.append(img)
    #         bboxes.append(box)
    #     images = np.array(images)
    #     bboxes = np.array(bboxes)
    #     return images, bboxes


if __name__ == "__main__":
    Dataset()
